{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (Efficient) Spherical Harmonic Transform\n",
    "\n",
    "## Preliminaries\n",
    "\n",
    "### Spherical harmonics transform\n",
    "\n",
    "The truncated series expansion of a function $f$ defined on the surface of a sphere can be written as\n",
    "\n",
    "$$\n",
    "f(\\theta, \\lambda) = \\sum_{m=-M}^{M} \\exp(im\\lambda) \\sum_{n=|m|}^{M} F_n^m \\bar{P}_n^m (\\cos \\theta),\n",
    "$$\n",
    "\n",
    "where $\\theta$ is the colatitude, $\\lambda$ the longitude, $\\bar{P}_n^m$ the normalized, associated Legendre polynomials and $F_n^m$, the expansion coefficient associated to the mode $(m,n)$.\n",
    "\n",
    "A direct spherical harmonic transform can be accomplished by a Fourier transform\n",
    "\n",
    "$$\n",
    "F^m(\\theta) = \\frac{1}{2 \\pi} \\int_{0}^{2\\pi} f(\\theta, \\lambda) \\exp(-im\\lambda) \\; \\mathrm{d}\\lambda\n",
    "$$\n",
    "\n",
    "in longitude and a Legendre transform\n",
    "\n",
    "$$\n",
    "F_n^m = \\frac{1}{2} \\int_{-1}^1 F^m(\\theta) \\bar{P}_n^m(\\cos \\theta) \\; \\mathrm{d} \\cos \\theta\n",
    "$$\n",
    "\n",
    "in latitude."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Usage\n",
    "\n",
    "### Preparation\n",
    "\n",
    "Let us begin by loading some libraries that we will be needing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import paddle\n",
    "import paddle.nn as nn\n",
    "import paddle.fft as fft\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "cmap = 'turbo'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We should also load the `paddle-sht`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from paddle_harmonics import *"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If a GPU is available we want to use it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = paddle.device.set_device('gpu' if paddle.device.cuda.device_count() > 0 else 'cpu')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test data\n",
    "\n",
    "Let us load some test-data in 2D for testing. We use the MOLA digital elevation map of MARS, courtesy of NASA:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p ./data\n",
    "!wget https://astropedia.astrogeology.usgs.gov/download/Mars/GlobalSurveyor/MOLA/thumbs/Mars_MGS_MOLA_DEM_mosaic_global_1024.jpg -O ./data/mola_topo.jpg"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us load and visualize the data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import imageio.v3 as iio\n",
    "\n",
    "img = iio.imread('./data/mola_topo.jpg')\n",
    "#convert to grayscale\n",
    "data = np.dot(img[...,:3]/255, [0.299, 0.587, 0.114])\n",
    "# interpolate onto 512x1024 grid:\n",
    "data = nn.functional.interpolate(paddle.to_tensor(data).unsqueeze(0).unsqueeze(0), size=(512,1024)).squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_theta = data.shape[0]\n",
    "n_lambda = data.shape[1]\n",
    "lon = np.linspace(-np.pi, np.pi, n_lambda)\n",
    "lat = np.linspace(np.pi/2., -np.pi/2., n_theta)\n",
    "Lon, Lat = np.meshgrid(lon, lat)\n",
    "\n",
    "fig = plt.figure(figsize=(10, 5))\n",
    "ax = fig.add_subplot(1, 1, 1, projection='mollweide')\n",
    "ax.pcolormesh(Lon, Lat, data, cmap=cmap)\n",
    "# ax.set_title(\"Elevation map of mars\")\n",
    "ax.grid(True)\n",
    "ax.set_xticklabels([])\n",
    "ax.set_yticklabels([])\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We want the data on the GPU. Moreover, `paddle_harmonics` expects the last two dimensions to be the spatial dimensions. To this end, we permute the data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "signal = data.to(device)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Spherical Harmonic Transform:\n",
    "\n",
    "To allow for composability, `paddle_harmonics` provides it's transform in the form of Pypaddle NN modules. As such we begin by initializing them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sht = RealSHT(n_theta, n_lambda, grid=\"equiangular\").to(device)\n",
    "isht = InverseRealSHT(n_theta, n_lambda, grid=\"equiangular\").to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "coeffs = sht(signal)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us plot the logarithm of the absolute coefficients:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(22, 5))\n",
    "ax = fig.add_subplot(1, 2, 1)\n",
    "plt.imshow(coeffs.abs().log().numpy(), cmap=cmap)\n",
    "plt.title(\"$\\log | SHT(u) |$\")\n",
    "plt.colorbar()\n",
    "\n",
    "ax = fig.add_subplot(1, 2, 2, projection='mollweide')\n",
    "plt.pcolormesh(Lon, Lat, isht(coeffs).cpu(), cmap=cmap)\n",
    "ax.set_title(\"Reconstructed elevation map\")\n",
    "ax.grid(True)\n",
    "ax.set_xticklabels([])\n",
    "ax.set_yticklabels([])\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can transform the the solution back and plot the error:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(10, 5))\n",
    "ax = fig.add_subplot(1, 1, 1, projection='mollweide')\n",
    "img = ax.pcolormesh(Lon, Lat, (isht(coeffs).numpy() - data).abs(), cmap=cmap, vmin=0, vmax=0.01)\n",
    "ax.set_title(\"Interpolation error\")\n",
    "ax.grid(True)\n",
    "ax.set_xticklabels([])\n",
    "ax.set_yticklabels([])\n",
    "plt.colorbar(img, ax=ax)\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Differentiability\n",
    "\n",
    "`paddle_harmonics` is differentiable, so let's use this functionality. We can perform a simple fit of spectral coefficients by using a simple model which will use 100 modes to approximate our original model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SpectralModel(nn.Layer):\n",
    "    def __init__(self, n_modes, out_dims):\n",
    "        super().__init__()\n",
    "        self.coeffs = paddle.create_parameter([n_modes, n_modes+1, 2], dtype=paddle.float64, \n",
    "                                              default_initializer=paddle.nn.initializer.Normal(mean=0, std=1))\n",
    "        self.isht = InverseRealSHT(out_dims[0], out_dims[1], lmax=n_modes, mmax=n_modes+1, grid=\"equiangular\").to(device)\n",
    "\n",
    "    def forward(self):\n",
    "        return self.isht(self.coeffs.as_complex())\n",
    "\n",
    "sh_model = SpectralModel(120, (n_theta, n_lambda)).to(device)\n",
    "sh_model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and now, let us train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = paddle.optimizer.Adam(parameters=sh_model.parameters(), learning_rate = 5e-2)\n",
    "\n",
    "for iter in range(500):\n",
    "\n",
    "    loss = (sh_model() - signal).pow(2).mean()\n",
    "\n",
    "    optimizer.clear_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    if iter % 10 == 0:\n",
    "        print(f'iteration: {iter} loss: {loss.item()}')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "let us plot the approximation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(10, 5))\n",
    "ax = fig.add_subplot(1, 1, 1, projection='mollweide')\n",
    "ax.pcolormesh(Lon, Lat, sh_model().detach().cpu(), cmap=cmap)\n",
    "# ax.set_title(\"Elevation map of mars\")\n",
    "ax.grid(True)\n",
    "ax.set_xticklabels([])\n",
    "ax.set_yticklabels([])\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Not bad! Let us again plot the error:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(10, 5))\n",
    "ax = fig.add_subplot(1, 1, 1, projection='mollweide')\n",
    "img = ax.pcolormesh(Lon, Lat, (sh_model().detach() - signal).abs().cpu(), cmap=cmap, vmin=0, vmax=1)\n",
    "ax.set_title(\"Approximation error\")\n",
    "ax.grid(True)\n",
    "ax.set_xticklabels([])\n",
    "ax.set_yticklabels([])\n",
    "plt.colorbar(img, ax=ax)\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "This concludes our introduction to `paddle-harmonics`. If you are enjoying this library, make sure to cite us:"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
